# httprelay.py - simple HTTP relay service for Google's AppEngine
# Licensed under the MIT License:
#
# Copyright (c) 2010 Peter Arrenbrecht
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

from google.appengine.ext import webapp
from google.appengine.api import memcache
from google.appengine.ext.webapp.util import run_wsgi_app
import time, base64, pickle, logging


MAX_CHUNK_SIZE = 1000000 - 1024
# 1 MB minus some elbow room for AppSpot's memcache

MAX_DELAY = 1


class Cache(object):
    
    def __init__(self, request):
        self.namespace = request.host.split(".", 1)[0]

    def _get(self, key):
        return memcache.get(key, namespace=self.namespace)
    def _set(self, key, val, time):
        return memcache.set(key, val, namespace=self.namespace)
    def _incr(self, key, initial=0):
        return memcache.incr(key, namespace=self.namespace, initial_value=initial)
    def _del(self, key):
        return memcache.delete(key, namespace=self.namespace)

    def nextseqno(self, key):
        return str(self._incr(key))
    def currseqno(self, key):
        return str(int(self._get(key) or "0") + 1)
    def resetseqno(self, key):
        return self._set(key, 0, time=0)

    def _packetkey(self, basekey, chunk):
        return "%sc%i" % (basekey, chunk)

    def setpacket(self, key, val, time):
        pickled = pickle.dumps(val)
        total = len(pickled)
        chunks = (total + MAX_CHUNK_SIZE - 1) / MAX_CHUNK_SIZE
        nextkey = self._packetkey(key, 0)
        nextoffs = 0
        for c in xrange(chunks):
            thiskey, thisoffs = nextkey, nextoffs
            nextkey = self._packetkey(key, c + 1)
            nextoffs += MAX_CHUNK_SIZE
            self._del(nextkey) # ensure we see no stale data
            self._set(thiskey, [chunks, pickled[thisoffs:nextoffs]], time)

    def droppacket(self, key):
        return self._del(self._packetkey(key, 0))

    def _pollchunk(self, key):
        logging.info("Reading %s:%s" % (self.namespace, key))
        val = self._get(key)
        if val is None:
            logging.debug("Polling %s:%s" % (self.namespace, key))
            delay = 0.1
            while val is None:
                time.sleep(delay)
                val = self._get(key)
                if delay < MAX_DELAY:
                    delay = min(delay * 2, MAX_DELAY) # backoff
        self._del(key) # remove this when debugging
        return val

    def pollpacket(self, key):
        nchunks, chunk = self._pollchunk(self._packetkey(key, 0))
        if not nchunks:
            return pickle.loads(chunk)
        chunks = [chunk]
        for c in xrange(1, nchunks):
            nchunks, chunk = self._pollchunk(self._packetkey(key, c))
            chunks.append(chunk)
        return pickle.loads("".join(chunks))


def sendback(cache, key, response):
    headers, body = cache.pollpacket(key)
    logging.debug("Got response %s:%s" % (cache.namespace, key))
    for h in headers.splitlines():
        if h:
            k, v = h.split(": ", 1)
            response.headers[k] = v
    type = response.headers.get("Content-Type").lower()
    if not type or not type.startswith("text/"):
        body = base64.b64decode(body)
    response.out.write(body)


class ClientReq(webapp.RequestHandler):

    def get(self):
        return self.handle("GET")

    def put(self):
        return self.handle("POST")

    def post(self):
        return self.handle("POST")

    def handle(self, method):
        cache = Cache(self.request)
        path = self.request.path_qs
        headers = "".join("%s: %s\n" % (k,v) for k,v in self.request.headers.iteritems())
        body = str(self.request.body)
        id = cache.nextseqno("head")
        logging.info("%s:%s %s %s" % (cache.namespace, id, method, path))
        req = [method, path, headers, body]
        cache.droppacket("r" + id) # ensure we see no stale data
        cache.setpacket("q" + id, req, time=60)
        if cache.namespace == "local":
            self.response.headers["Content-Type"] = "text/plain"
            self.response.out.write("OK: %s" % id)
        else:
            sendback(cache, "r" + id, self.response)


class ClientFetchReq(webapp.RequestHandler):

    def get(self, id):
        sendback(Cache(self.request), "r" + id, self.response)


class ServerRestartReq(webapp.RequestHandler):

    def get(self):
        cache = Cache(self.request)
        # ensure poll sees no initial stale data
        cache.droppacket("q1")
        # reset request queue
        cache.resetseqno("head")
        cache.resetseqno("tail")


class ServerPollReq(webapp.RequestHandler):

    def get(self):
        cache = Cache(self.request)
        id = cache.currseqno("tail")
        method, path, headers, body = cache.pollpacket("q" + id)
        self.response.headers["Content-Type"] = "text/plain"
        self.response.out.write("httprelay %s %s %s\n" % (id, method, path))
        self.response.out.write(headers)
        self.response.out.write("\n-\n")
        self.response.out.write(body)
        cache.nextseqno("tail")


class ServerFwdReq(webapp.RequestHandler):

    def post(self, id):
        cache = Cache(self.request)
        logging.info("Got forward: %s:%s" % (cache.namespace, id))
        rsp = [self.request.get(p) for p in ("headers", "body")]
        cache.setpacket("r" + id, rsp, time=5000)
        self.response.headers["Content-Type"] = "text/plain"
        self.response.out.write("OK")


class PeekReq(webapp.RequestHandler):

    def get(self, key):
        self.response.headers["Content-Type"] = "text/plain"
        self.response.out.write("%s" % Cache(self.request)._get(key))


class ResetReq(webapp.RequestHandler):

    def get(self):
        self.response.headers["Content-Type"] = "text/plain"
        stats = memcache.get_stats()
        for k, v in stats.iteritems():
            self.response.out.write("%s: %s\n" % (k, v))
        memcache.flush_all()


class UploadReq(webapp.RequestHandler):

    def post(self):
        for k in ("Content-Type"):
            self.response.headers[k] = self.request.headers.get(k)
        self.response.out.write(len(self.request.body))


class DownloadReq(webapp.RequestHandler):

    def get(self, size):
        self.response.headers["Content-Type"] = "text/plain"
        for _i in xrange(int(size)):
            self.response.out.write("The quick brown fox jumps over the lazy dog.\n")


class NullReq(webapp.RequestHandler):

    def get(self):
        self.response.set_status(404)


def main():
    application = webapp.WSGIApplication(
        [
         (r"/httprelay/restart", ServerRestartReq),
         (r"/httprelay/poll", ServerPollReq),
         (r"/httprelay/fwd/(.*)", ServerFwdReq),
         (r"/httprelay/fetch/(.*)", ClientFetchReq),
         (r"/httprelay/peek/(.*)", PeekReq), # FIXME
         (r"/httprelay/reset", ResetReq), # FIXME
#         (r"/httprelay/upload", UploadReq), # FIXME
#         (r"/httprelay/download/(.*)", DownloadReq), # FIXME
         (r".*", ClientReq),
        ],
        debug=True)
    run_wsgi_app(application)

if __name__ == "__main__":
    main()
